from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED
from tqdm import tqdm


class ThreadPoolParams:
    def __init__(self, target, *args, **kwargs):
        self.target = target
        self.args = args
        self.kwargs = kwargs


class FixThreadPoolExecutor(object):
    """
    固定大小线程池执行器
    继承并重写generate_task()即可
    """

    def __init__(self, pool_size=2, task_total=0, desc="执行进度"):
        self.pool_size = pool_size
        self.tasks = []
        self.task_total = task_total
        self.desc = desc

    def run(self, task_iterable_func):
        """
        线程池执行
        :return:
        """
        tasks = []
        with ThreadPoolExecutor(max_workers=self.pool_size) as pool:
            with tqdm(total=self.task_total, desc=self.desc) as pbar:
                for task in task_iterable_func():
                    if type(task) != ThreadPoolParams:
                        raise Exception("type of task is not ThreadPoolParams")
                    info = pool.submit(task.target, *task.args, **task.kwargs)
                    tasks.append(info)
                    if len(tasks) > self.pool_size:
                        print(f"task gt {self.pool_size},please wait ...")
                        wait(tasks, return_when=ALL_COMPLETED)
                        pbar.update(len(tasks))
                        tasks = []
                if tasks:
                    wait(tasks, return_when=ALL_COMPLETED)
                    pbar.update(len(tasks))
                    tasks = []
